package chapter4.eg3;

import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Properties;
import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.errors.WakeupException;
import org.junit.Before;
import org.junit.Test;

public class RebalanceHandlerApp {

  private Map<TopicPartition, OffsetAndMetadata> currentOffsets = new HashMap<>();

  private KafkaConsumer<String, String> consumer;

  @Before
  public void before() {
    Properties properties = new Properties();
    properties.put("bootstrap.servers", "localhost:9092");
    properties.put("group.id", "CountryCounter");
    properties.put("key.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
    properties.put("value.deserializer", "org.apache.kafka.common.serialization.StringDeserializer");
    consumer = new KafkaConsumer<>(properties);
    consumer.subscribe(Collections.singletonList("customerCountries"), new HandleRebalance());
  }

  @Test
  public void testRebalanceCallback() {
    try {
      while (true) {
        ConsumerRecords<String, String> records = consumer.poll(100);
        for (ConsumerRecord<String, String> record : records) {
          System.out.printf("topic: %s, partition: %s, offset: %d, customer: %s, country: %s\n",
              record.topic(), record.partition(), record.offset(), record.key(), record.value());
          currentOffsets.put(new TopicPartition(record.topic(), record.partition()),
              new OffsetAndMetadata(record.offset() + 1, "no metadata"));
        }

        consumer.commitAsync(currentOffsets, null);
      }
    } catch (WakeupException e) {
      // ignore
    } catch (Exception e) {
      e.printStackTrace();
    } finally {
      try {
        consumer.commitSync(currentOffsets);
      } finally {
        consumer.close();
        System.out.println("Closed consumer and we are done");
      }
    }
  }

  private final class HandleRebalance implements ConsumerRebalanceListener {

    @Override
    public void onPartitionsRevoked(Collection<TopicPartition> partitions) {
      System.out.println("Lost partitions in rebalance. Committing current offsets: " + currentOffsets);
      consumer.commitSync(currentOffsets);
    }

    @Override
    public void onPartitionsAssigned(Collection<TopicPartition> partitions) {

    }
  }


}
